# -*- coding: utf-8 -*-

"""
图像生成器
"""

from PIL import Image, ImageDraw, ImageFont
import random
import os
import numpy as np


class ImageGenerator:
    def __init__(self):
        """初始化图像生成器"""
        pass
        
    def generate_puzzle_image(self, theme, size=(800, 600)):
        """
        生成拼图图像
        
        Args:
            theme (str): 图像主题
            size (tuple): 图像尺寸
            
        Returns:
            PIL.Image: 生成的图像
        """
        # 创建新图像
        img = Image.new('RGB', size, self._get_background_color(theme))
        draw = ImageDraw.Draw(img)
        
        # 根据主题绘制元素
        if "风景" in theme:
            self._draw_landscape(draw, size, theme)
        elif "城市" in theme:
            self._draw_city(draw, size, theme)
        elif "抽象" in theme:
            self._draw_abstract(draw, size, theme)
        else:
            self._draw_generic(draw, size, theme)
            
        return img
        
    def _get_background_color(self, theme):
        """根据主题获取背景色"""
        color_map = {
            "阳光明媚": (135, 206, 235),  # 天蓝色
            "雨景": (70, 70, 70),         # 灰色
            "山水": (34, 139, 34),        # 森林绿
            "城市": (25, 25, 112),        # 午夜蓝
            "抽象": (128, 128, 128),      # 灰色
            "自然": (34, 139, 34)         # 森林绿
        }
        
        for key, color in color_map.items():
            if key in theme:
                return color
        return (200, 200, 200)  # 默认灰色
        
    def _draw_landscape(self, draw, size, theme):
        """绘制风景"""
        width, height = size
        
        # 绘制地面
        ground_y = height * 0.7
        draw.rectangle([0, ground_y, width, height], fill=(34, 100, 34))
        
        # 绘制太阳
        if "阳光" in theme:
            draw.ellipse([width-100, 50, width-50, 100], fill=(255, 255, 0))
            
        # 绘制云朵
        for _ in range(3):
            x = random.randint(0, width-100)
            y = random.randint(50, int(ground_y)-100)
            draw.ellipse([x, y, x+80, y+40], fill=(255, 255, 255))
            
        # 绘制树
        for _ in range(5):
            x = random.randint(50, width-50)
            draw.rectangle([x, ground_y-80, x+20, ground_y], fill=(101, 67, 33))
            draw.ellipse([x-30, ground_y-120, x+50, ground_y-40], fill=(34, 139, 34))
            
    def _draw_city(self, draw, size, theme):
        """绘制城市"""
        width, height = size
                                  
        # 绘制建筑物
        building_count = random.randint(5, 10)
        for i in range(building_count):
            building_width = random.randint(60, 120)
            building_height = random.randint(100, height // 2)
            x = i * (width // building_count)
            draw.rectangle(
                [x, height - building_height, x + building_width, height],
                fill=(random.randint(50, 150), random.randint(50, 150), random.randint(50, 150))
            )
                                      
        # 绘制窗户
        for i in range(building_count):
            building_width = random.randint(60, 120)
            building_height = random.randint(100, height // 2)
            x = i * (width // building_count)
            for j in range(5):
                for k in range(5):
                    if random.random() > 0.3:  # 70%概率绘制窗户
                        wx = x + 10 + j * 15
                        wy = height - building_height + 20 + k * 20
                        draw.rectangle([wx, wy, wx + 10, wy + 15], fill=(255, 255, 0))

    def _draw_abstract(self, draw, size, theme):
        """绘制抽象图案"""
        width, height = size
        
        # 绘制随机形状
        for _ in range(20):
            shape_type = random.choice(["circle", "rectangle", "line"])
            color = (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255))
            
            if shape_type == "circle":
                x = random.randint(0, width)
                y = random.randint(0, height)
                r = random.randint(10, 50)
                draw.ellipse([x-r, y-r, x+r, y+r], fill=color)
            elif shape_type == "rectangle":
                x = random.randint(0, width)
                y = random.randint(0, height)
                w = random.randint(20, 100)
                h = random.randint(20, 100)
                draw.rectangle([x, y, x+w, y+h], fill=color)
            else:  # line
                x1 = random.randint(0, width)
                y1 = random.randint(0, height)
                x2 = random.randint(0, width)
                y2 = random.randint(0, height)
                draw.line([x1, y1, x2, y2], fill=color, width=random.randint(1, 5))
                
    def _draw_generic(self, draw, size, theme):
        """绘制通用图案"""
        width, height = size
        
        # 绘制渐变背景
        for i in range(height):
            r = int(255 * (i / height))
            g = int(200 * (i / height))
            b = int(150 * (i / height))
            draw.line([(0, i), (width, i)], fill=(r, g, b))
            
        # 绘制中心图案
        center_x, center_y = width // 2, height // 2
        for i in range(5):
            radius = 100 - i * 20
            draw.ellipse([center_x-radius, center_y-radius, 
                         center_x+radius, center_y+radius], 
                        outline=(255, 255, 255), width=2)

    def generate_image_with_stable_diffusion(self, prompt, style, resolution=(1024, 1024)):
        """
        使用Stable Diffusion XL Turbo生成图像
        
        Args:
            prompt (str): 图像主题词
            style (str): 艺术风格
            resolution (tuple): 图像分辨率 (width, height)
            
        Returns:
            PIL.Image: 生成的图像
        """
        try:
            # 导入必要的库
            import torch
            from diffusers import AutoPipelineForImage2Image
            from PIL import Image
            import io
            
            # 检查是否有GPU可用
            device = "cuda" if torch.cuda.is_available() else "cpu"
            print(f"Using device: {device}")
            
            # 加载Stable Diffusion XL Turbo管道
            pipe = AutoPipelineForImage2Image.from_pretrained(
                "stabilityai/sdxl-turbo",
                torch_dtype=torch.float16 if device == "cuda" else torch.float32,
                variant="fp16" if device == "cuda" else None
            )
            pipe = pipe.to(device)
            
            # 创建初始图像（纯色背景）
            init_image = Image.new("RGB", (512, 512), (255, 255, 255))
            
            # 构造完整的提示词
            full_prompt = f"{prompt}, {style} style"
            
            # 生成图像
            image = pipe(
                prompt=full_prompt,
                image=init_image,
                num_inference_steps=2,  # SDXL Turbo只需要很少的步骤
                strength=0.9,  # 控制图像变化程度
                guidance_scale=0.0  # SDXL Turbo推荐值
            ).images[0]
            
            # 调整图像大小到指定分辨率
            image = image.resize(resolution, Image.LANCZOS)
            
            return image
            
        except ImportError:
            print("警告: 未安装diffusers库，使用默认图像生成方法")
            return self.generate_puzzle_image(f"{prompt} {style}", resolution)
        except Exception as e:
            print(f"图像生成出错: {e}")
            return self.generate_puzzle_image(f"{prompt} {style}", resolution)

    def split_image_to_puzzle_pieces(self, image, grid_size=3):
        """
        将图像分割为拼图碎片
        
        Args:
            image (PIL.Image): 输入图像
            grid_size (int): 网格大小 (3表示3x3, 4表示4x4等)
            
        Returns:
            list: 包含拼图碎片的列表，每个元素为(dict, PIL.Image)元组
        """
        width, height = image.size
        piece_width = width // grid_size
        piece_height = height // grid_size
        
        pieces = []
        
        # 为每个拼图碎片生成边缘匹配数据
        edge_profiles = self._generate_edge_profiles(grid_size)
        
        for row in range(grid_size):
            for col in range(grid_size):
                # 计算当前碎片的边界
                left = col * piece_width
                upper = row * piece_height
                right = left + piece_width
                lower = upper + piece_height
                
                # 添加一些重叠以确保边缘匹配
                overlap = 10  # 重叠像素数
                
                # 调整边界以包含重叠区域
                if col > 0:  # 不是第一列
                    left -= overlap
                if col < grid_size - 1:  # 不是最后一列
                    right += overlap
                if row > 0:  # 不是第一行
                    upper -= overlap
                if row < grid_size - 1:  # 不是最后一行
                    lower += overlap
                
                # 裁剪图像
                piece_image = image.crop((left, upper, right, lower))
                
                # 创建碎片信息
                piece_info = {
                    "position": (row, col),
                    "grid_size": grid_size,
                    "size": (piece_width, piece_height),
                    "edges": edge_profiles[row][col],
                    "original_position": (row, col)  # 用于游戏逻辑
                }
                
                pieces.append((piece_info, piece_image))
                
        return pieces

    def _generate_edge_profiles(self, grid_size):
        """
        生成边缘匹配数据以确保几何逻辑自洽
        
        Args:
            grid_size (int): 网格大小
            
        Returns:
            list: 二维列表，包含每个碎片的边缘信息
        """
        # 初始化边缘配置
        profiles = [[{} for _ in range(grid_size)] for _ in range(grid_size)]
        
        # 为每个碎片生成边缘类型
        for row in range(grid_size):
            for col in range(grid_size):
                # 上边缘
                if row == 0:
                    profiles[row][col]["top"] = "flat"
                else:
                    # 与上方碎片的下边缘匹配
                    profiles[row][col]["top"] = self._get_complementary_edge(
                        profiles[row-1][col]["bottom"]
                    ) if "bottom" in profiles[row-1][col] else "out"
                
                # 左边缘
                if col == 0:
                    profiles[row][col]["left"] = "flat"
                else:
                    # 与左方碎片的右边缘匹配
                    profiles[row][col]["left"] = self._get_complementary_edge(
                        profiles[row][col-1]["right"]
                    ) if "right" in profiles[row][col-1] else "out"
                
                # 下边缘
                if row == grid_size - 1:
                    profiles[row][col]["bottom"] = "flat"
                else:
                    # 随机生成，但确保能与相邻碎片匹配
                    profiles[row][col]["bottom"] = random.choice(["out", "in"])
                
                # 右边缘
                if col == grid_size - 1:
                    profiles[row][col]["right"] = "flat"
                else:
                    # 随机生成，但确保能与相邻碎片匹配
                    profiles[row][col]["right"] = random.choice(["out", "in"])
        
        return profiles

    def _get_complementary_edge(self, edge_type):
        """
        获取互补的边缘类型
        
        Args:
            edge_type (str): 原始边缘类型
            
        Returns:
            str: 互补的边缘类型
        """
        if edge_type == "out":
            return "in"
        elif edge_type == "in":
            return "out"
        else:
            return edge_type
